# renew_one_p_o.py
from datasets.hico_text_label import *

# annotated_actions  = ['ride', 'sit_on', 'row']
# hoi_labels = {'file_name': 'HICO_train2015_00001924.jpg', 'img_id': 1924, 'annotations': [{'bbox': [230, 248, 330, 331], 'category_id': 1}, {'bbox': [202, 316, 364, 366], 'category_id': 9}], 'hoi_annotation': [{'subject_id': 0, 'object_id': 1, 'category_id': 44, 'hoi_category_id': 36}, {'subject_id': 0, 'object_id': 1, 'category_id': 77, 'hoi_category_id': 39}]}

# true_hoi_labels = {'file_name': 'HICO_train2015_00001924.jpg', 'img_id': 1924, 'annotations': [{'bbox': [229, 236, 323, 332], 'category_id': 1}, {'bbox': [203, 314, 371, 360], 'category_id': 9}, {'bbox': [219, 238, 331, 329], 'category_id': 1}, {'bbox': [185, 303, 362, 373], 'category_id': 9}], 'hoi_annotation': [{'subject_id': 0, 'object_id': 1, 'category_id': 77, 'hoi_category_id': 39}, {'subject_id': 0, 'object_id': 1, 'category_id': 88, 'hoi_category_id': 42}, {'subject_id': 2, 'object_id': 3, 'category_id': 78, 'hoi_category_id': 40}]}
# print(true_hoi_labels)


# 排除人的干扰 找到物体的id
def find_object_category_id(hoi_labels):
    object_category_ids = []

    # 获取 hoi_labels 中的 annotations
    annotations = hoi_labels["annotations"]

    # 遍历 annotations，提取 category_id
    for annotation in annotations:
        # 获取 category_id  这里要注意
        category_id = valid_obj_ids.index(annotation["category_id"])
        print(get_object_name(category_id))
        # 如果 category_id 不是 1（代表人类），则将其添加到列表中
        if category_id != 0:
            object_category_ids.append(category_id)

    return object_category_ids


# 获取物体的名称 可以再改 遍历读取id 再遍历输出名字
# 这里是一个人一个物体 就只需要取出第一个 先暂时这样写
def get_object_name(object_category_id):
    # 获取物体类别的ID
    object_name = obj2name.get(object_category_id, "Unknown")
    return object_name


# 找到物体id之后 从先验条件里找到物体对应的动作范围
def get_object_actions(hoi_labels):
    # 获取物体类别的ID
    object_category_ids = find_object_category_id(hoi_labels)
    object_category_id = object_category_ids[0]
    # object_category_id_min_1 = object_category_id - 1

    # 使用 hico_object2label_index 字典获取对应的动作标签索引范围
    action_label_range = hico_object2label_index.get(object_category_id)

    if action_label_range:
        action_label_range_add_1 = [x + 1 for x in action_label_range]
        start_index, end_index = action_label_range_add_1
        print(
            f"Action label indices for object category {object_category_id}: {start_index} to {end_index}"
        )
        # 提取动作的名字并打印输出
        object_actions = [
            hoi_index2action[i] for i in range(start_index, end_index + 1)
        ]
        print(
            f"Actions for object category {object_category_id} (name {get_object_name(object_category_id)}): {object_actions }"
        )
    else:
        print(f"No action labels found for object category {object_category_id}")

    return start_index, object_actions


# 只调用这个函数 这个函数在调用这个py文件里的其他函数 算作是main函数
def renew_hoi_labels(annotated_actions, hoi_labels):
    # 在这里实现后处理逻辑，根据人工标记的动作信息和框架生成的 HOI 标签
    # 可以根据需要修改 hoi_labels
    print("renew_one_p_o.py")
    print("后处理前标签：")
    print(hoi_labels)
    # 从原先已有的hoi标签的物体信息获取动作的起始和截止索引 是根据
    start_index, object_actions = get_object_actions(hoi_labels)
    renewed_hoi_labels = hoi_labels.copy()
    # 判断hoi_annotation中的hoi_category_id是否在action_label_range范围内，不在则删除
    # 但其实这个没用 是一定在的一开始就有mask了的 可能多个物体会有用吧 好像确实没用
    # 重要的是把没有的给加上去
    # 不在人工标注标签的先删掉 也先好代码 后面再问师兄需不需要删掉
    # renewed_hoi_labels['hoi_annotation'] = [hoi_anno for hoi_anno in hoi_labels['hoi_annotation']
    #                                         if action_label_range[0] <= hoi_anno['hoi_category_id']
    #                                         <= action_label_range[1]]

    # 循环处理每个annotated action 人工注释的动作
    for annotated_action in annotated_actions:
        # 复制一个hoi_annotation的条目并保存为临时变量  这里应该需要改 一个人一个物体这里是空的 需要先最初有个初始的样子 可以复制一个 也可以随意构造一个乱初始化的
        temp_hoi_anno = hoi_labels["hoi_annotation"][0].copy()

        # 根据annotated action获取新的category_id
        new_category_id = action2index.get(annotated_action)
        # 更新temp_hoi_anno中的category_id和hoi_category_id
        temp_hoi_anno["category_id"] = new_category_id

        # 这里会有报错
        # 查找annotated_action在object_actions中的位置
        # 不能直接这样写 忽略了目标检测失误的情况
        # action_index = object_actions.index(annotated_action)

        # 如果 annotated_action 不在 object_actions 中，可以选择跳过当前循环
        # 就说明目标检测漏检测了 人工标注的数据没能用上 只能丢掉这个标签
        if annotated_action in object_actions:
            action_index = object_actions.index(annotated_action)
        else:
            continue

        # 根据action_index和action2index获取新的hoi_category_id
        new_category_id = action_index + start_index  # 假设已经有了start_index
        temp_hoi_anno["hoi_category_id"] = new_category_id

        # 将修改后的temp_hoi_anno添加到renewed_hoi_labels中
        renewed_hoi_labels["hoi_annotation"].append(temp_hoi_anno)

    return renewed_hoi_labels


# renewed_hoi_labels = renew_hoi_labels(annotated_actions, hoi_labels)
# print(renewed_hoi_labels)
